package db

import (
	"database/sql"
	"fmt"
	"sort"
	"sync"

	"mcp-db-gateway/pkg/logger"
	"mcp-db-gateway/pkg/storage"

	"github.com/google/uuid"
	_ "github.com/go-sql-driver/mysql"
	_ "github.com/lib/pq"
	"go.uber.org/zap"
)

type Manager struct {
	mu      sync.RWMutex
	conns   map[string]*sql.DB
	txs     map[string]*sql.Tx // Active transactions, sessionID -> tx
	aliases []string
	store   *storage.Store
}

func NewManager(store *storage.Store) (*Manager, error) {
	m := &Manager{
		conns: make(map[string]*sql.DB),
		txs:   make(map[string]*sql.Tx),
		store: store,
	}

	storedConnections, err := store.GetAllConnections()
	if err != nil {
		return nil, fmt.Errorf("failed to get connections from store: %w", err)
	}

	for _, dbConfig := range storedConnections {
		db, err := sql.Open(dbConfig.Type, dbConfig.DSN)
		if err != nil {
			// Log error but continue, so one bad config doesn't stop everything
			logger.Log.Warn("Failed to open database",
				zap.String("alias", dbConfig.Alias),
				zap.Error(err),
			)
			continue
		}

		if err := db.Ping(); err != nil {
			logger.Log.Warn("Failed to ping database",
				zap.String("alias", dbConfig.Alias),
				zap.Error(err),
			)
			db.Close() // close the connection if ping fails
			continue
		}

		m.conns[dbConfig.Alias] = db
		m.aliases = append(m.aliases, dbConfig.Alias)
	}
	sort.Strings(m.aliases)

	return m, nil
}

func (m *Manager) AddConnection(alias, dbType, dsn string) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	// Check if alias already exists
	if _, exists := m.conns[alias]; exists {
		return fmt.Errorf("database alias '%s' already exists", alias)
	}

	// Persist to store first
	if err := m.store.AddConnection(alias, dbType, dsn); err != nil {
		return fmt.Errorf("failed to save connection to store: %w", err)
	}

	// Then, add to the in-memory pool
	db, err := sql.Open(dbType, dsn)
	if err != nil {
		// If it fails, we should ideally roll back the store operation
		// For now, we remove it to keep consistency.
		m.store.RemoveConnection(alias)
		return fmt.Errorf("failed to open new database connection: %w", err)
	}
	if err := db.Ping(); err != nil {
		m.store.RemoveConnection(alias)
		return fmt.Errorf("failed to ping new database connection: %w", err)
	}

	m.conns[alias] = db
	m.aliases = append(m.aliases, alias)
	sort.Strings(m.aliases)

	return nil
}

func (m *Manager) RemoveConnection(alias string) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	// Check if alias exists
	db, exists := m.conns[alias]
	if !exists {
		return fmt.Errorf("database alias '%s' not found", alias)
	}

	// Remove from store first
	if err := m.store.RemoveConnection(alias); err != nil {
		return fmt.Errorf("failed to remove connection from store: %w", err)
	}

	// Then, close and remove from the in-memory pool
	db.Close()
	delete(m.conns, alias)

	// Rebuild aliases slice
	var newAliases []string
	for i, a := range m.aliases {
		if a == alias {
			newAliases = append(m.aliases[:i], m.aliases[i+1:]...)
			break
		}
	}
	m.aliases = newAliases

	return nil
}

// BeginTx starts a new transaction for the given database alias.
// It returns a unique session ID for the transaction.
func (m *Manager) BeginTx(alias string) (string, error) {
	m.mu.Lock()
	defer m.mu.Unlock()

	db, ok := m.conns[alias]
	if !ok {
		return "", fmt.Errorf("database alias '%s' not found", alias)
	}

	tx, err := db.Begin()
	if err != nil {
		return "", err
	}

	sessionID := uuid.New().String()
	m.txs[sessionID] = tx

	return sessionID, nil
}

// GetTx retrieves an active transaction by its session ID.
func (m *Manager) GetTx(sessionID string) (*sql.Tx, bool) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	tx, ok := m.txs[sessionID]
	return tx, ok
}

// CommitTx commits a transaction and removes it from the manager.
func (m *Manager) CommitTx(sessionID string) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	tx, ok := m.txs[sessionID]
	if !ok {
		return fmt.Errorf("transaction with session ID '%s' not found", sessionID)
	}

	if err := tx.Commit(); err != nil {
		return err
	}
	delete(m.txs, sessionID)
	return nil
}

// RollbackTx rolls back a transaction and removes it from the manager.
func (m *Manager) RollbackTx(sessionID string) error {
	m.mu.Lock()
	defer m.mu.Unlock()

	tx, ok := m.txs[sessionID]
	if !ok {
		return fmt.Errorf("transaction with session ID '%s' not found", sessionID)
	}

	if err := tx.Rollback(); err != nil {
		return err
	}
	delete(m.txs, sessionID)
	return nil
}

// HealthCheck iterates over all connections and pings them, returning a map of their status.
func (m *Manager) HealthCheck() map[string]string {
	m.mu.RLock()
	defer m.mu.RUnlock()

	status := make(map[string]string)
	for alias, db := range m.conns {
		if err := db.Ping(); err != nil {
			status[alias] = "FAILED: " + err.Error()
		} else {
			status[alias] = "OK"
		}
	}
	return status
}

func (m *Manager) GetDB(alias string) (*sql.DB, bool) {
	m.mu.RLock()
	defer m.mu.RUnlock()
	db, ok := m.conns[alias]
	return db, ok
}

func (m *Manager) GetDBAliases() []string {
	m.mu.RLock()
	defer m.mu.RUnlock()
	// Return a copy to prevent modification of the internal slice
	aliasesCopy := make([]string, len(m.aliases))
	copy(aliasesCopy, m.aliases)
	return aliasesCopy
}

func (m *Manager) Close() {
	m.mu.Lock()
	defer m.mu.Unlock()
	for _, db := range m.conns {
		db.Close()
	}
} 